{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Trip planning with a FalkorDB GraphRAG agent using GroupChat\n",
    "\n",
    "In this notebook, we're building a trip planning system using the new AG2 GroupChat functionality which has an objective to create an itinerary together with a customer. The end result will be an itinerary that has route times and distances calculated between activities.\n",
    "\n",
    "The following diagram outlines the key components of the GroupChat, with highlights being:\n",
    "\n",
    "- FalkorDB agent using a GraphRAG database of restaurants and attractions\n",
    "- Structured Output agent that will enforce a strict format for the accepted itinerary\n",
    "- Routing agent that utilises the Google Maps API to calculate distances between activities\n",
    "- GroupChat orchestration utilising context variables and handoff conditions\n",
    "\n",
    "**Note:** This notebook has been updated from the deprecated Swarm functionality to use the new GroupChat approach as described in the [migration guide](https://docs.ag2.ai/0.9.8/docs/user-guide/advanced-concepts/orchestration/swarm/deprecation/)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "````mdx-code-block\n",
    ":::note\n",
    "This notebook has been updated from the deprecated Swarm functionality to use the new GroupChat approach. The GroupChat system now accommodates any ConversableAgent and provides better orchestration through handoff conditions.\n",
    ":::\n",
    "````"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "````{=mdx}\n",
    ":::info Requirements\n",
    "FalkorDB's GraphRAG-SDK is a dependency for this notebook, which can be installed with ag2 via pip:\n",
    "\n",
    "```bash\n",
    "pip install -U ag2[openai]\n",
    "pip install graphrag_sdk\n",
    "```\n",
    "\n",
    "> **Note:** If you have been using `autogen` or `ag2`, all you need to do is upgrade it using:  \n",
    "> ```bash\n",
    "> pip install -U autogen[openai]\n",
    "> ```\n",
    "> or  \n",
    "> ```bash\n",
    "> pip install -U ag2[openai]\n",
    "> ```\n",
    "> as `autogen`, and `ag2` are aliases for the same PyPI package.  \n",
    "\n",
    "\n",
    "For more information, please refer to the [installation guide](https://docs.ag2.ai/latest/docs/user-guide/basic-concepts/installing-ag2).\n",
    ":::\n",
    "````\n",
    "\n",
    "### Pydantic\n",
    "\n",
    "Please ensure you have Pydantic version 2+ installed.\n",
    "```bash\n",
    "pip install --upgrade pydantic\n",
    "```\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Running a FalkorDB\n",
    "\n",
    "**Note:** You need to have a FalkorDB graph database running. If you are running one in a Docker container, please ensure your Docker network is setup to allow access to it.\n",
    "\n",
    "In this example, we've set the FalkorDB host and port, please adjust them accordingly. For how to set up FalkorDB, please refer to [https://docs.falkordb.com/](https://docs.falkordb.com/).\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Google Maps API Key\n",
    "\n",
    "To use Google's API to calculate travel times, you will need to have enabled the `Directions API` in your Google Maps Platform. You can get an API key and free quota, see [here](https://developers.google.com/maps/documentation/directions/overview) and [here](https://developers.google.com/maps/get-started) for more details.\n",
    "\n",
    "Once you have your API key, set your environment variable `GOOGLE_MAP_API_KEY` to the key"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set Configuration and OpenAI API Key\n",
    "\n",
    "By default, FalkorDB uses OpenAI LLMs and that requires an OpenAI key in your environment variable `OPENAI_API_KEY`.\n",
    "\n",
    "You can utilise an OAI_CONFIG_LIST file and extract the OpenAI API key and put it in the environment, as will be shown in the following cell.\n",
    "\n",
    "Alternatively, you can load the environment variable yourself.\n",
    "\n",
    "````{=mdx}\n",
    ":::tip\n",
    "Learn more about configuring LLMs for agents [here](https://docs.ag2.ai/latest/docs/user-guide/basic-concepts/llm-configuration).\n",
    ":::\n",
    "````"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import autogen\n",
    "\n",
    "llm_config = autogen.LLMConfig.from_json(path=\"OAI_CONFIG_LIST\", timeout=120).where(model=\"gpt-5-nano\")\n",
    "\n",
    "# Put the OpenAI API key into the environment\n",
    "# os.environ[\"OPENAI_API_KEY\"] = llm_config.config_list[0].api_key"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prepare the FalkorDB GraphRAG database\n",
    "\n",
    "Using 3 sample JSON data files from our GitHub repository, we will create a specific ontology for our GraphRAG database and then populate it.\n",
    "\n",
    "Creating a specific ontology that matches with the types of queries makes for a more optimal database and is more cost efficient when populating the knowledge graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from autogen.agentchat.contrib.graph_rag.document import Document, DocumentType\n",
    "\n",
    "# 3 Files (adjust path as necessary)\n",
    "input_paths = [\n",
    "    \"../test/agentchat/contrib/graph_rag/trip_planner_data/attractions.jsonl\",\n",
    "    \"../test/agentchat/contrib/graph_rag/trip_planner_data/cities.jsonl\",\n",
    "    \"../test/agentchat/contrib/graph_rag/trip_planner_data/restaurants.jsonl\",\n",
    "]\n",
    "input_documents = [Document(doctype=DocumentType.TEXT, path_or_url=input_path) for input_path in input_paths]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create Ontology\n",
    "\n",
    "Entities: Country, City, Attraction, Restaurant\n",
    "\n",
    "Relationships: City in Country, Attraction in City, Restaurant in City"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from graphrag_sdk import Attribute, AttributeType, Entity, Ontology, Relation\n",
    "\n",
    "# Attraction + Restaurant + City + Country Ontology\n",
    "trip_data_ontology = Ontology()\n",
    "\n",
    "trip_data_ontology.add_entity(\n",
    "    Entity(\n",
    "        label=\"Country\",\n",
    "        attributes=[\n",
    "            Attribute(\n",
    "                name=\"name\",\n",
    "                attr_type=AttributeType.STRING,\n",
    "                required=True,\n",
    "                unique=True,\n",
    "            ),\n",
    "        ],\n",
    "    )\n",
    ")\n",
    "trip_data_ontology.add_entity(\n",
    "    Entity(\n",
    "        label=\"City\",\n",
    "        attributes=[\n",
    "            Attribute(\n",
    "                name=\"name\",\n",
    "                attr_type=AttributeType.STRING,\n",
    "                required=True,\n",
    "                unique=True,\n",
    "            ),\n",
    "            Attribute(\n",
    "                name=\"weather\",\n",
    "                attr_type=AttributeType.STRING,\n",
    "                required=False,\n",
    "                unique=False,\n",
    "            ),\n",
    "            Attribute(\n",
    "                name=\"population\",\n",
    "                attr_type=AttributeType.NUMBER,\n",
    "                required=False,\n",
    "                unique=False,\n",
    "            ),\n",
    "        ],\n",
    "    )\n",
    ")\n",
    "trip_data_ontology.add_entity(\n",
    "    Entity(\n",
    "        label=\"Restaurant\",\n",
    "        attributes=[\n",
    "            Attribute(\n",
    "                name=\"name\",\n",
    "                attr_type=AttributeType.STRING,\n",
    "                required=True,\n",
    "                unique=True,\n",
    "            ),\n",
    "            Attribute(\n",
    "                name=\"description\",\n",
    "                attr_type=AttributeType.STRING,\n",
    "                required=False,\n",
    "                unique=False,\n",
    "            ),\n",
    "            Attribute(\n",
    "                name=\"rating\",\n",
    "                attr_type=AttributeType.NUMBER,\n",
    "                required=False,\n",
    "                unique=False,\n",
    "            ),\n",
    "            Attribute(\n",
    "                name=\"food_type\",\n",
    "                attr_type=AttributeType.STRING,\n",
    "                required=False,\n",
    "                unique=False,\n",
    "            ),\n",
    "        ],\n",
    "    )\n",
    ")\n",
    "trip_data_ontology.add_entity(\n",
    "    Entity(\n",
    "        label=\"Attraction\",\n",
    "        attributes=[\n",
    "            Attribute(\n",
    "                name=\"name\",\n",
    "                attr_type=AttributeType.STRING,\n",
    "                required=True,\n",
    "                unique=True,\n",
    "            ),\n",
    "            Attribute(\n",
    "                name=\"description\",\n",
    "                attr_type=AttributeType.STRING,\n",
    "                required=False,\n",
    "                unique=False,\n",
    "            ),\n",
    "            Attribute(\n",
    "                name=\"type\",\n",
    "                attr_type=AttributeType.STRING,\n",
    "                required=False,\n",
    "                unique=False,\n",
    "            ),\n",
    "        ],\n",
    "    )\n",
    ")\n",
    "trip_data_ontology.add_relation(\n",
    "    Relation(\n",
    "        label=\"IN_COUNTRY\",\n",
    "        source=\"City\",\n",
    "        target=\"Country\",\n",
    "    )\n",
    ")\n",
    "trip_data_ontology.add_relation(\n",
    "    Relation(\n",
    "        label=\"IN_CITY\",\n",
    "        source=\"Restaurant\",\n",
    "        target=\"City\",\n",
    "    )\n",
    ")\n",
    "trip_data_ontology.add_relation(\n",
    "    Relation(\n",
    "        label=\"IN_CITY\",\n",
    "        source=\"Attraction\",\n",
    "        target=\"City\",\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Establish FalkorDB and load\n",
    "\n",
    "Remember: Change your host, port, and preferred OpenAI model if needed (gpt-5-nano and better is recommended)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from graphrag_sdk.models.openai import OpenAiGenerativeModel\n",
    "\n",
    "from autogen.agentchat.contrib.graph_rag.falkor_graph_query_engine import FalkorGraphQueryEngine\n",
    "from autogen.agentchat.contrib.graph_rag.falkor_graph_rag_capability import FalkorGraphRagCapability\n",
    "\n",
    "# Create FalkorGraphQueryEngine\n",
    "query_engine = FalkorGraphQueryEngine(\n",
    "    name=\"trip_data\",\n",
    "    # host=\"192.168.65.1\",  # Change\n",
    "    port=6379,  # if needed\n",
    "    ontology=trip_data_ontology,\n",
    "    model=OpenAiGenerativeModel(\"gpt-5-nano\"),\n",
    ")\n",
    "\n",
    "# Ingest data and initialize the database\n",
    "# query_engine.init_db(input_doc=input_documents)\n",
    "\n",
    "# If you have already ingested and created the database, you can use this connect_db instead of init_db\n",
    "query_engine.connect_db()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# IMPORTS\n",
    "import json\n",
    "\n",
    "import requests\n",
    "from pydantic import BaseModel\n",
    "\n",
    "from autogen import ConversableAgent, UserProxyAgent\n",
    "from autogen.agentchat import initiate_group_chat\n",
    "from autogen.agentchat.group import (\n",
    "    AgentTarget,\n",
    "    ContextVariables,\n",
    "    OnCondition,\n",
    "    ReplyResult,\n",
    "    RevertToUserTarget,\n",
    "    StringLLMCondition,\n",
    ")\n",
    "from autogen.agentchat.group.patterns import AutoPattern\n",
    "from autogen.agentchat.group.targets.transition_target import AgentNameTarget"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Pydantic model for Structured Output\n",
    "\n",
    "Utilising OpenAI's [Structured Outputs](https://platform.openai.com/docs/guides/structured-outputs), our Structured Output agent's responses will be constrained to this Pydantic model.\n",
    "\n",
    "The itinerary is structured as:\n",
    "Itinerary has Day(s) has Event(s)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Event(BaseModel):\n",
    "    type: str  # Attraction, Restaurant, Travel\n",
    "    location: str\n",
    "    city: str\n",
    "    description: str\n",
    "\n",
    "\n",
    "class Day(BaseModel):\n",
    "    events: list[Event]\n",
    "\n",
    "\n",
    "class Itinerary(BaseModel):\n",
    "    days: list[Day]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Google Maps Platform\n",
    "\n",
    "The functions necessary to query the Directions API to get travel times."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _fetch_travel_time(origin: str, destination: str) -> dict:\n",
    "    \"\"\"Retrieves route information using Google Maps Directions API.\n",
    "    API documentation at https://developers.google.com/maps/documentation/directions/get-directions\n",
    "    \"\"\"\n",
    "    endpoint = \"https://maps.googleapis.com/maps/api/directions/json\"\n",
    "    params = {\n",
    "        \"origin\": origin,\n",
    "        \"destination\": destination,\n",
    "        \"mode\": \"walking\",  # driving (default), bicycling, transit\n",
    "        \"key\": os.environ.get(\"GOOGLE_MAP_API_KEY\"),\n",
    "    }\n",
    "\n",
    "    response = requests.get(endpoint, params=params)\n",
    "    if response.status_code == 200:\n",
    "        return response.json()\n",
    "    else:\n",
    "        return {\"error\": \"Failed to retrieve the route information\", \"status_code\": response.status_code}\n",
    "\n",
    "\n",
    "def update_itinerary_with_travel_times(context_variables: ContextVariables) -> ReplyResult:\n",
    "    \"\"\"Update the complete itinerary with travel times between each event.\"\"\"\n",
    "    \"\"\"\n",
    "    Retrieves route information using Google Maps Directions API.\n",
    "    API documentation at https://developers.google.com/maps/documentation/directions/get-directions\n",
    "    \"\"\"\n",
    "\n",
    "    # Ensure that we have a structured itinerary, if not, back to the structured_output_agent to make it\n",
    "    if context_variables.get(\"structured_itinerary\") is None:\n",
    "        print(\"DEBUG: No structured_itinerary found, redirecting to structured_output_agent\")\n",
    "        return ReplyResult(\n",
    "            target=AgentNameTarget(\"structured_output_agent\"),\n",
    "            message=\"Structured itinerary not found, please create the structured output, structured_output_agent.\",\n",
    "        )\n",
    "    elif \"timed_itinerary\" in context_variables:\n",
    "        print(\"DEBUG: Timed itinerary already exists, task complete\")\n",
    "        # Return to user since work is done\n",
    "        return ReplyResult(\n",
    "            target=RevertToUserTarget(),\n",
    "            context_variables=context_variables,\n",
    "            message=\"Timed itinerary already done, inform the customer that their itinerary is ready!\",\n",
    "        )\n",
    "\n",
    "    # Process the itinerary, converting it back to an object and working through each event to work out travel time and distance\n",
    "    itinerary_object = Itinerary.model_validate(json.loads(context_variables[\"structured_itinerary\"]))\n",
    "    for day in itinerary_object.days:\n",
    "        events = day.events\n",
    "        new_events = []\n",
    "        pre_event, cur_event = None, None\n",
    "        event_count = len(events)\n",
    "        index = 0\n",
    "        while index < event_count:\n",
    "            if index > 0:\n",
    "                pre_event = events[index - 1]\n",
    "\n",
    "            cur_event = events[index]\n",
    "            if pre_event:\n",
    "                origin = \", \".join([pre_event.location, pre_event.city])\n",
    "                destination = \", \".join([cur_event.location, cur_event.city])\n",
    "                maps_api_response = _fetch_travel_time(origin=origin, destination=destination)\n",
    "                try:\n",
    "                    leg = maps_api_response[\"routes\"][0][\"legs\"][0]\n",
    "                    travel_time_txt = f\"{leg['duration']['text']}, ({leg['distance']['text']})\"\n",
    "                    new_events.append(\n",
    "                        Event(\n",
    "                            type=\"Travel\",\n",
    "                            location=f\"walking from {pre_event.location} to {cur_event.location}\",\n",
    "                            city=cur_event.city,\n",
    "                            description=travel_time_txt,\n",
    "                        )\n",
    "                    )\n",
    "                except Exception:\n",
    "                    print(f\"Note: Unable to get travel time from {origin} to {destination}\")\n",
    "            new_events.append(cur_event)\n",
    "            index += 1\n",
    "        day.events = new_events\n",
    "\n",
    "    context_variables[\"timed_itinerary\"] = itinerary_object.model_dump()\n",
    "\n",
    "    # Task complete - return to user with final context\n",
    "    return ReplyResult(\n",
    "        target=RevertToUserTarget(),\n",
    "        context_variables=context_variables,\n",
    "        message=\"Timed itinerary added to context with travel times. Your itinerary is ready!\",\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## GroupChat\n",
    "\n",
    "### Context Variables\n",
    "Our GroupChat agents will have access to a couple of context variables in relation to the itinerary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trip_context = ContextVariables({\n",
    "    \"itinerary_confirmed\": False,\n",
    "    \"itinerary\": \"\",\n",
    "    \"structured_itinerary\": None,\n",
    "})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Agent Functions\n",
    "\n",
    "We have two functions/tools for our agents.\n",
    "\n",
    "One for our Planner agent to mark an itinerary as confirmed by the customer and to store the final text itinerary. This will then transfer to our Structured Output agent.\n",
    "\n",
    "Another for the Structured Output Agent to save the structured itinerary and transfer to the Route Timing agent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Agent function placeholder - will be defined after agents are created\n",
    "def mark_itinerary_as_complete(final_itinerary: str, context_variables: ContextVariables) -> ReplyResult:\n",
    "    \"\"\"Store and mark our itinerary as accepted by the customer.\"\"\"\n",
    "\n",
    "    context_variables[\"itinerary_confirmed\"] = True\n",
    "    context_variables[\"itinerary\"] = final_itinerary\n",
    "\n",
    "    # This will update the context variables and then transfer to the Structured Output agent\n",
    "    return ReplyResult(\n",
    "        target=AgentNameTarget(\"structured_output_agent\"),\n",
    "        context_variables=context_variables,\n",
    "        message=\"Itinerary recorded and confirmed.\",\n",
    "    )\n",
    "\n",
    "\n",
    "def create_structured_itinerary(context_variables: ContextVariables, structured_itinerary: str) -> ReplyResult:\n",
    "    \"\"\"Once a structured itinerary is created, store it and pass on to the Route Timing agent.\"\"\"\n",
    "\n",
    "    # Store the structured itinerary regardless of confirmation status\n",
    "    # The confirmation should happen at the planner level, not here\n",
    "    context_variables[\"structured_itinerary\"] = structured_itinerary\n",
    "\n",
    "    # This will update the context variables and then transfer to the Route Timing agent\n",
    "    return ReplyResult(\n",
    "        target=AgentNameTarget(\"route_timing_agent\"),\n",
    "        context_variables=context_variables,\n",
    "        message=\"Structured itinerary stored.\",\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Agents\n",
    "\n",
    "Our GroupChat agents and a UserProxyAgent (human) which the group will interact with."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Planner agent, interacting with the customer and GraphRag agent, to create an itinerary\n",
    "planner_agent = ConversableAgent(\n",
    "    name=\"planner_agent\",\n",
    "    system_message=\"You are a trip planner agent. It is important to know where the customer is going, how many days, what they want to do.\"\n",
    "    + \"You will work with another agent, graphrag_agent, to get information about restaurant and attractions. \"\n",
    "    + \"You are also working with the customer, so you must ask the customer what they want to do if you don't have LOCATION, NUMBER OF DAYS, MEALS, and ATTRACTIONS. \"\n",
    "    + \"When you have the customer's requirements, work with graphrag_agent to get information for an itinerary.\"\n",
    "    + \"You are responsible for creating the itinerary and for each day in the itinerary you MUST HAVE events and EACH EVENT MUST HAVE a 'type' ('Restaurant' or 'Attraction'), 'location' (name of restaurant or attraction), 'city', and 'description'. \"\n",
    "    + \"Finally, YOU MUST ask the customer if they are happy with the itinerary before marking the itinerary as complete.\",\n",
    "    functions=[mark_itinerary_as_complete],\n",
    "    llm_config=llm_config,\n",
    ")\n",
    "\n",
    "# FalkorDB GraphRAG agent, utilising the FalkorDB to gather data for the Planner agent\n",
    "graphrag_agent = ConversableAgent(\n",
    "    name=\"graphrag_agent\",\n",
    "    system_message=\"Return a list of restaurants and/or attractions. List them separately and provide ALL the options in the location. Do not provide travel advice.\",\n",
    ")\n",
    "\n",
    "# Adding the FalkorDB capability to the agent\n",
    "graph_rag_capability = FalkorGraphRagCapability(query_engine)\n",
    "graph_rag_capability.add_to_agent(graphrag_agent)\n",
    "\n",
    "# Structured Output agent, formatting the itinerary into a structured format through the response_format on the LLM Configuration\n",
    "structured_llm_config = autogen.LLMConfig.from_json(path=\"OAI_CONFIG_LIST\", timeout=120).where(model=\"gpt-5-nano\")\n",
    "\n",
    "for config in structured_llm_config.config_list:\n",
    "    config.response_format = Itinerary\n",
    "\n",
    "structured_output_agent = ConversableAgent(\n",
    "    name=\"structured_output_agent\",\n",
    "    system_message=\"You are a data formatting agent, format the provided itinerary in the context below into the provided format. IMPORTANT: You must call the create_structured_itinerary function to save the formatted itinerary.\",\n",
    "    llm_config=structured_llm_config,\n",
    "    functions=[create_structured_itinerary],\n",
    ")\n",
    "\n",
    "# Route Timing agent, adding estimated travel times to the itinerary by utilising the Google Maps Platform\n",
    "route_timing_agent = ConversableAgent(\n",
    "    name=\"route_timing_agent\",\n",
    "    system_message=\"You are a route timing agent. Your job is to call the update_itinerary_with_travel_times tool to add travel times to the itinerary. The tool will handle everything including ending the conversation when done.\",\n",
    "    llm_config=llm_config,\n",
    "    functions=[update_itinerary_with_travel_times],\n",
    ")\n",
    "\n",
    "# Our customer will be a human in the loop\n",
    "customer = UserProxyAgent(name=\"customer\")\n",
    "\n",
    "# Set up handoffs for each agent using the proper handoffs methods\n",
    "# NOTE: For handoffs, we use AgentTarget with actual agent objects\n",
    "planner_agent.handoffs.add_llm_conditions([\n",
    "    OnCondition(\n",
    "        target=AgentTarget(graphrag_agent),\n",
    "        condition=StringLLMCondition(\"Need information on the restaurants and attractions for a location\"),\n",
    "    ),\n",
    "    OnCondition(\n",
    "        target=AgentTarget(structured_output_agent),\n",
    "        condition=StringLLMCondition(\"Itinerary is confirmed by the customer\"),\n",
    "    ),\n",
    "])\n",
    "\n",
    "# Back to the Planner when information has been retrieved\n",
    "graphrag_agent.handoffs.set_after_work(AgentTarget(planner_agent))\n",
    "\n",
    "# Once we have formatted our itinerary, we can hand off to the route timing agent to add in the travel timings\n",
    "structured_output_agent.handoffs.set_after_work(AgentTarget(route_timing_agent))\n",
    "\n",
    "# Route timing agent will revert to user when done (handled by the tool function)\n",
    "# Note: No explicit handoff needed since the tool handles returning to user"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### GroupChat Pattern Setup\n",
    "\n",
    "Instead of the old Swarm handoff system, we now use GroupChat patterns with handoff conditions directly on agents. We'll use the AutoPattern which dynamically selects the next speaker based on conversation context.\n",
    "\n",
    "For more details on the new GroupChat orchestration, [see the documentation](https://docs.ag2.ai/0.9.8/docs/user-guide/advanced-concepts/orchestration/group-chat/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create AutoPattern for dynamic agent selection\n",
    "pattern = AutoPattern(\n",
    "    initial_agent=planner_agent,\n",
    "    agents=[planner_agent, graphrag_agent, structured_output_agent, route_timing_agent],\n",
    "    user_agent=customer,\n",
    "    context_variables=trip_context,\n",
    "    group_manager_args={\"llm_config\": llm_config},\n",
    "    # group_after_work=RevertToUserTarget(),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run the GroupChat\n",
    "\n",
    "Let's get an itinerary for a couple of days in Rome using the AutoPattern."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start the conversation using GroupChat with AutoPattern\n",
    "chat_result, final_context, last_agent = initiate_group_chat(\n",
    "    pattern=pattern,\n",
    "    messages=\"I want to go to Rome for a couple of days. Can you help me plan my trip?\",\n",
    "    max_rounds=100,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Bonus itinerary output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_itinerary(itinerary_data):\n",
    "    header = \"█             █\\\\n █           █ \\\\n  █  █████  █  \\\\n   ██     ██   \\\\n  █         █  \\\\n █  ███████  █ \\\\n █ ██ ███ ██ █ \\\\n   █████████   \\\\n\\\\n ██   ███ ███  \\\\n█  █ █       █ \\\\n████ █ ██  ██  \\\\n█  █ █  █ █    \\\\n█  █  ██  ████ \\\\n\"\n",
    "    width = 80\n",
    "    icons = {\"Travel\": \"🚶\", \"Restaurant\": \"🍽️\", \"Attraction\": \"🏛️\"}\n",
    "\n",
    "    for line in header.split(\"\\\\n\"):\n",
    "        print(line.center(width))\n",
    "    print(f\"Itinerary for {itinerary_data['days'][0]['events'][0]['city']}\".center(width))\n",
    "    print(\"=\" * width)\n",
    "\n",
    "    for day_num, day in enumerate(itinerary_data[\"days\"], 1):\n",
    "        print(f\"\\\\nDay {day_num}\".center(width))\n",
    "        print(\"-\" * width)\n",
    "\n",
    "        for event in day[\"events\"]:\n",
    "            event_type = event[\"type\"]\n",
    "            print(f\"\\\\n  {icons[event_type]} {event['location']}\")\n",
    "            if event_type != \"Travel\":\n",
    "                words = event[\"description\"].split()\n",
    "                line = \"    \"\n",
    "                for word in words:\n",
    "                    if len(line) + len(word) + 1 <= 76:\n",
    "                        line += word + \" \"\n",
    "                    else:\n",
    "                        print(line)\n",
    "                        line = \"    \" + word + \" \"\n",
    "                if line.strip():\n",
    "                    print(line)\n",
    "            else:\n",
    "                print(f\"    {event['description']}\")\n",
    "        print(\"\\\\n\" + \"-\" * width)\n",
    "\n",
    "\n",
    "# First, try to get the itinerary from context variables\n",
    "print(\"Checking context variables for itinerary data...\")\n",
    "\n",
    "# Check multiple potential sources for the itinerary\n",
    "itinerary_data = None\n",
    "itinerary_found = False\n",
    "\n",
    "# 1. Check final_context\n",
    "if final_context.get(\"timed_itinerary\"):\n",
    "    print(\"Found timed_itinerary in final_context!\")\n",
    "    itinerary_data = final_context[\"timed_itinerary\"]\n",
    "    itinerary_found = True\n",
    "elif final_context.get(\"structured_itinerary\"):\n",
    "    print(\"Found structured_itinerary in final_context!\")\n",
    "    itinerary_data = json.loads(final_context[\"structured_itinerary\"])\n",
    "    itinerary_found = True\n",
    "\n",
    "# 2. Check pattern context variables if final_context didn't have it\n",
    "if not itinerary_found and hasattr(pattern, \"context_variables\"):\n",
    "    pattern_context = pattern.context_variables\n",
    "    if pattern_context.get(\"timed_itinerary\"):\n",
    "        print(\"Found timed_itinerary in pattern context!\")\n",
    "        itinerary_data = pattern_context[\"timed_itinerary\"]\n",
    "        itinerary_found = True\n",
    "    elif pattern_context.get(\"structured_itinerary\"):\n",
    "        print(\"Found structured_itinerary in pattern context!\")\n",
    "        itinerary_data = json.loads(pattern_context[\"structured_itinerary\"])\n",
    "        itinerary_found = True\n",
    "\n",
    "# 3. If context variables failed, extract from chat messages as fallback\n",
    "if not itinerary_found:\n",
    "    print(\"Context variables empty, extracting itinerary from chat messages...\")\n",
    "\n",
    "    # Look through the chat result messages for structured output\n",
    "    if hasattr(chat_result, \"chat_history\"):\n",
    "        messages = chat_result.chat_history\n",
    "    elif hasattr(chat_result, \"messages\"):\n",
    "        messages = chat_result.messages\n",
    "    else:\n",
    "        messages = []\n",
    "\n",
    "    print(f\"Checking {len(messages)} messages for itinerary data...\")\n",
    "\n",
    "    for i, message in enumerate(messages):\n",
    "        try:\n",
    "            # Check if this message contains structured itinerary JSON\n",
    "            content = message.get(\"content\", \"\") if isinstance(message, dict) else str(message)\n",
    "\n",
    "            # Look for JSON-like content that matches our Itinerary structure\n",
    "            if \"days\" in content and \"events\" in content and \"type\" in content:\n",
    "                print(f\"Found potential itinerary in message {i}\")\n",
    "\n",
    "                # Try to extract JSON from the content\n",
    "                import re\n",
    "\n",
    "                json_match = re.search(r'\\\\{.*\"days\".*\\\\}', content, re.DOTALL)\n",
    "                if json_match:\n",
    "                    try:\n",
    "                        itinerary_json = json_match.group(0)\n",
    "                        itinerary_data = json.loads(itinerary_json)\n",
    "                        print(\"Successfully parsed itinerary JSON from messages!\")\n",
    "                        itinerary_found = True\n",
    "                        break\n",
    "                    except json.JSONDecodeError:\n",
    "                        continue\n",
    "        except Exception:\n",
    "            continue\n",
    "\n",
    "# Display the itinerary if found\n",
    "if itinerary_found and itinerary_data:\n",
    "    print(\"\\\\n\" + \"=\" * 80)\n",
    "    print(\"ITINERARY FOUND! Printing...\")\n",
    "    print(\"=\" * 80)\n",
    "    print_itinerary(itinerary_data)\n",
    "else:\n",
    "    print(\"No valid itinerary found anywhere.\")\n",
    "    print(\"\\\\nDEBUG INFO:\")\n",
    "    print(\"Final context:\", final_context.to_dict() if hasattr(final_context, \"to_dict\") else final_context)\n",
    "    print(\n",
    "        \"Pattern context:\",\n",
    "        pattern.context_variables.to_dict()\n",
    "        if hasattr(pattern, \"context_variables\") and hasattr(pattern.context_variables, \"to_dict\")\n",
    "        else \"No pattern context\",\n",
    "    )\n",
    "\n",
    "    # Show available keys in contexts\n",
    "    if hasattr(final_context, \"keys\") or hasattr(final_context, \"to_dict\"):\n",
    "        available_keys = (\n",
    "            list(final_context.keys())\n",
    "            if hasattr(final_context, \"keys\")\n",
    "            else list(final_context.to_dict().keys())\n",
    "            if hasattr(final_context, \"to_dict\")\n",
    "            else \"Unknown\"\n",
    "        )\n",
    "        print(f\"Available keys in final_context: {available_keys}\")\n",
    "\n",
    "    if hasattr(pattern, \"context_variables\"):\n",
    "        pattern_keys = (\n",
    "            list(pattern.context_variables.keys())\n",
    "            if hasattr(pattern.context_variables, \"keys\")\n",
    "            else list(pattern.context_variables.to_dict().keys())\n",
    "            if hasattr(pattern.context_variables, \"to_dict\")\n",
    "            else \"Unknown\"\n",
    "        )\n",
    "        print(f\"Available keys in pattern context: {pattern_keys}\")"
   ]
  }
 ],
 "metadata": {
  "front_matter": {
   "description": "FalkorDB GraphRAG utilises a knowledge graph and can be added as a capability to agents. Together with a swarm orchestration of agents is highly effective at providing a RAG capability.",
   "tags": [
    "RAG",
    "tool/function",
    "swarm"
   ]
  },
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
